package blackRedTree;


public class RBTree {
    //内部类
    class RBTreeNode {
        int val;
        Enum color;
        RBTreeNode left;
        RBTreeNode right;
        RBTreeNode parent;

        public RBTreeNode(int val) {
            this.val = val;
            this.color = Color.RED;
        }
    }

    RBTreeNode root;

    public boolean add(int val) {
        if (root == null) {
            root = new RBTreeNode(val);
            root.color = Color.BLACK;
            return true;
        }
        RBTreeNode parent = null;
        RBTreeNode cur = root;
        while (cur != null) {//寻找插入位置
            parent = cur;
            if (val > cur.val) {
                cur = cur.right;
            } else if (val == cur.val) {
                return false;
            } else {
                cur = cur.left;
            }
        }
        RBTreeNode node = new RBTreeNode(val);
        if (parent.val > val) {
            parent.left = node;
        } else {
            parent.right = node;
        }
        node.parent = parent;
        cur = node;


        //插入完成后看是否有两重复的红色parent.parent != null
        while (parent != null && parent.color == Color.RED) {
            //左边出现两个连续红色
            RBTreeNode g = parent.parent;
            if(g.left == parent){
                RBTreeNode u = g.right;
                if(u != null && u.color == Color.RED){  //情况一: cur为红，p为红，g为黑，u存在且为红
                    parent.color = Color.BLACK;
                    u.color = Color.BLACK;
                    g.color = Color.RED;
                    cur = g;
                    parent = cur.parent;
                }else if(cur.color == Color.RED && g.color == Color.BLACK){
//                    情况三: cur为红，p为红，g为黑，u不存在/u为黑
                    if(parent.right == cur){
                        leftRevolve(cur);
                        RBTreeNode tep = cur;
                        cur = parent;
                        parent = tep;
                    }
                    //情况二: cur为红，p为红，g为黑，u不存在/u为黑
                    //右旋
                    rightRevolve(parent);
                    parent.color = Color.BLACK;
                    g.color = Color.RED;
                }

            }else if(g.right == parent){
                //右边出现两个连续红色
                RBTreeNode u = g.left;
                if(u != null && u.color == Color.RED){  //情况一: cur为红，p为红，g为黑，u存在且为红
                    parent.color = Color.BLACK;
                    u.color = Color.BLACK;
                    g.color = Color.RED;
                    cur = g;
                    parent = cur.parent;
                }else if(cur.color == Color.RED && g.color == Color.BLACK){
//                    情况三: cur为红，p为红，g为黑，u不存在/u为黑
                    if(parent.left == cur){
                        rightRevolve(cur);
                        RBTreeNode tep = cur;
                        cur = parent;
                        parent = tep;
                    }
                    //情况二: cur为红，p为红，g为黑，u不存在/u为黑
                    //左旋
                    leftRevolve(parent);
                    parent.color = Color.BLACK;
                    g.color = Color.RED;
                }
            }
        }
        //调整完成后，根节点调为黑色
        root.color = Color.BLACK;
        return true;
    }

    //右单旋
    private void rightRevolve(RBTreeNode cur) {
        RBTreeNode parent = cur.parent;
        RBTreeNode curR = cur.right;
        RBTreeNode parentP = parent.parent;
        parent.left = curR;
        if(curR != null){
            curR.parent = parent;
        }
        cur.right = parent;
        parent.parent = cur;
        cur.parent = parentP;

        if(parentP == null){
            root = cur;
        }else {
            if(parentP.right == parent){
                parentP.right = cur;
            }else {
                parentP.left = cur;
            }
        }
    }

    //左单旋
    private void leftRevolve(RBTreeNode cur) {
        RBTreeNode parent = cur.parent;
        RBTreeNode curL = cur.left;
        RBTreeNode pParent = parent.parent;

        parent.right = curL;
        if(curL != null){
            curL.parent = parent;
        }
        cur.left = parent;
        parent.parent = cur;
        cur.parent = pParent;
        if(pParent == null){
            root = cur;
        }else {
            if(pParent.left == parent){
                pParent.left = cur;
            }else {
                pParent.right = cur;
            }
        }
    }
    public boolean check(){
        RBTreeNode cur = root;
        if(cur == null){
            return true;
        }
        if(cur.color != Color.BLACK){
            System.out.println("根结点不为黑色！");
            return false;
        }

        int count = 0;
        while (cur != null){
            if (cur.color == Color.BLACK){
                count++;
            }
            cur = cur.left;
        }

        return isHaveTwoRedNode(root) && isIdenticalBlackNodeSum(root,0,count);
    }

    private boolean isIdenticalBlackNodeSum(RBTreeNode r,int path,int count) {
        if(r == null){
            return true;
        }
        if(r.color == Color.BLACK){
            path++;
        }

        if(r.right == null && r.left == null){
            if(path == count){
                return true;
            }else {
                System.out.println("任意路径的黑色节点数目不同！");
                return false;
            }


        }

        return isIdenticalBlackNodeSum(r.left,path,count) && isIdenticalBlackNodeSum(r.right,path,count);
    }

    private boolean isHaveTwoRedNode(RBTreeNode r) {
        if(r == null){
            return true;
        }

        if(r.color == Color.RED){
            if(r.parent.color == Color.RED){
                System.out.println("有两个连续的红色节点！");
                return false;
            }
        }

        return isHaveTwoRedNode(r.left) && isHaveTwoRedNode(r.right);
    }

    //中序遍历
    public void inOrder(RBTreeNode r){
        if (r == null){
            return;
        }
        inOrder(r.left);
        System.out.print(r.val + " ");
        inOrder(r.right);
    }

}

